import pysrt
import json
import os
import chardet


def parse_subtitles(file_path):
    with open(file_path, 'rb') as f:
        # 自动检测文件编码
        raw_data = f.read()
        result = chardet.detect(raw_data)
        encoding = result['encoding']

    # 使用检测到的编码读取文件
    subs = pysrt.open(file_path, encoding=encoding)
    subtitles = []

    for sub in subs:
        timestamp = str(sub.start) + " --> " + str(sub.end)
        text = sub.text.replace("\n", " ").strip()

        file_name = os.path.basename(file_path)

        subtitle_entry = {
            "instruction": file_name,
            "input": timestamp,
            "output": text
        }
        subtitles.append(subtitle_entry)

    return subtitles


# 将字幕数据保存为 JSON 文件
def save_to_json(subtitles, output_path):
    with open(output_path, 'w', encoding='utf-8') as f:
        json.dump(subtitles, f, ensure_ascii=False, indent=4)


# 遍历目录中的所有 SRT 文件并提取内容
def extract_all_subtitles(input_dir, output_json_path):
    all_subtitles = []

    # 遍历目录中的所有文件
    for file_name in os.listdir(input_dir):
        if file_name.endswith('.srt'):  # 只处理 SRT 文件
            file_path = os.path.join(input_dir, file_name)
            subtitles = parse_subtitles(file_path)
            all_subtitles.extend(subtitles)  # 将当前文件的字幕内容添加到总列表中

    # 保存所有字幕到一个 JSON 文件
    save_to_json(all_subtitles, output_json_path)
    print(f"字幕已保存到 {output_json_path}")

y = 2022
for i in range(1, 13):
    # 示例：将所有字幕文件转换为指定格式并保存
    input_subtitle_dir = fr'D:\subtitle\moves\{y}\{i}月'  # 字幕文件夹路径

    if i < 10:
        output_json_path = fr'D:\subtitle\moves\json\{y}0{i}.json'  # 输出 JSON 文件路径
    else:
        output_json_path = fr'D:\subtitle\moves\json\{y}{i}.json'

    extract_all_subtitles(input_subtitle_dir, output_json_path)
